{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "388e8477-57af-4bd9-9ad4-ecef6ed162c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "将数据集转为yolo格式\n",
    "\n",
    "@Author: enpei\n",
    "@Date: 2023-08-02\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0a59d61",
   "metadata": {},
   "source": [
    "目标格式：\n",
    ". 工作路径\n",
    "├── datasets\n",
    "│   └── custom_dataset\n",
    "│       ├── images\n",
    "│       │   ├── train\n",
    "│       │   │   └── demo_001.jpg\n",
    "│       │   └── val\n",
    "│       │       └── demo_002.jpg\n",
    "│       └── labels\n",
    "│           ├── train\n",
    "│           │   └── demo_001.txt\n",
    "│           └── val\n",
    "│               └── demo_002.txt\n",
    "└── 其他文件、文件夹"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "56662226",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a34fd3e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "训练集图片数量： 6471\n",
      "训练集标注数量： 6471\n",
      "测试集图片数量： 548\n",
      "测试集标注数量： 548\n"
     ]
    }
   ],
   "source": [
    "# 训练集目录\n",
    "train_dir = \"./VisDrone2019-DET-train/\"\n",
    "# 测试集目录\n",
    "val_dir = \"./VisDrone2019-DET-val/\"\n",
    "\n",
    "# 训练集图片目录、标注目录\n",
    "train_img_dir = train_dir + \"/images\"\n",
    "train_anno_dir = train_dir + \"/annotations\"\n",
    "# 测试集图片目录、标注目录\n",
    "test_img_dir = val_dir + \"/images\"\n",
    "test_anno_dir = val_dir + \"/annotations\"\n",
    "\n",
    "# 统计数量\n",
    "print(\"训练集图片数量：\", len(os.listdir(train_img_dir)))\n",
    "print(\"训练集标注数量：\", len(os.listdir(train_anno_dir)))\n",
    "print(\"测试集图片数量：\", len(os.listdir(test_img_dir)))\n",
    "print(\"测试集标注数量：\", len(os.listdir(test_anno_dir)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e987f157",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "85c62631",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 查看一个标注sample\n",
    "sample_label_file = glob.glob(train_anno_dir + '/*.txt')[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "358f4753",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'./VisDrone2019-DET-train//annotations\\\\0000002_00005_d_0000014.txt'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sample_label_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d0d0f5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 读取内容\n",
    "with open(sample_label_file, \"r\", encoding=\"utf-8\") as fd:\n",
    "    lines = fd.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "59e10076",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['684,8,273,116,0,0,0,0\\n',\n",
       " '406,119,265,70,0,0,0,0\\n',\n",
       " '255,22,119,128,0,0,0,0\\n',\n",
       " '1,3,209,78,0,0,0,0\\n',\n",
       " '708,471,74,33,1,4,0,1\\n',\n",
       " '639,425,61,46,1,4,0,0\\n',\n",
       " '594,399,64,51,1,4,0,0\\n',\n",
       " '562,390,61,38,1,4,0,0\\n',\n",
       " '540,372,65,33,1,4,0,1\\n',\n",
       " '514,333,68,35,1,4,0,0\\n',\n",
       " '501,317,64,31,1,4,0,1\\n',\n",
       " '501,299,45,28,1,4,0,1\\n',\n",
       " '489,284,48,27,1,4,0,1\\n',\n",
       " '463,262,48,29,1,4,0,0\\n',\n",
       " '458,252,49,22,1,4,0,1\\n',\n",
       " '448,242,45,20,1,4,0,1\\n',\n",
       " '442,230,49,19,1,4,0,1\\n',\n",
       " '439,214,45,21,1,4,0,1\\n',\n",
       " '429,208,42,19,1,4,0,1\\n',\n",
       " '420,199,43,20,1,4,0,1\\n',\n",
       " '398,188,41,18,1,4,0,1\\n',\n",
       " '46,391,14,26,1,2,0,0\\n',\n",
       " '421,433,74,44,1,4,0,1\\n',\n",
       " '369,346,64,34,1,4,0,0\\n',\n",
       " '398,410,72,46,1,4,0,1\\n',\n",
       " '394,393,70,36,1,4,0,1\\n',\n",
       " '377,364,71,38,1,4,0,0\\n',\n",
       " '357,312,58,31,1,4,0,0\\n',\n",
       " '359,298,54,22,1,4,0,2\\n',\n",
       " '348,283,43,28,1,5,0,1\\n',\n",
       " '345,271,52,19,1,4,0,1\\n',\n",
       " '340,260,60,18,1,5,0,1\\n',\n",
       " '340,250,52,16,1,4,0,1\\n',\n",
       " '332,231,54,22,1,5,0,1\\n',\n",
       " '323,213,45,25,1,5,0,0\\n',\n",
       " '317,195,45,31,1,6,0,1\\n',\n",
       " '316,188,36,15,1,4,0,2\\n',\n",
       " '308,179,44,17,1,4,0,1\\n',\n",
       " '345,163,37,18,1,4,0,0\\n',\n",
       " '384,164,26,22,1,4,0,1\\n',\n",
       " '43,398,18,17,1,10,0,1\\n',\n",
       " '324,167,14,6,1,3,0,0\\n',\n",
       " '362,143,30,18,1,4,1,0\\n',\n",
       " '310,150,36,15,1,4,0,0\\n',\n",
       " '258,375,30,43,1,4,0,2\\n',\n",
       " '218,348,18,28,1,7,0,0\\n',\n",
       " '249,297,23,30,1,4,0,2\\n',\n",
       " '0,298,26,34,1,8,1,0\\n',\n",
       " '33,279,24,30,1,4,0,0\\n',\n",
       " '47,243,26,28,1,4,0,0\\n',\n",
       " '72,202,28,28,1,5,0,0\\n',\n",
       " '87,125,15,15,1,4,0,0\\n',\n",
       " '170,113,15,13,1,4,0,0\\n',\n",
       " '102,81,13,13,1,4,0,0\\n',\n",
       " '415,207,5,9,1,1,0,0\\n',\n",
       " '428,197,8,13,1,1,0,1\\n',\n",
       " '232,385,13,31,0,11,0,0\\n',\n",
       " '234,391,8,17,1,2,0,2\\n',\n",
       " '879,256,12,19,1,1,0,0\\n',\n",
       " '897,225,8,17,1,1,0,0\\n',\n",
       " '57,341,11,20,1,2,0,0\\n',\n",
       " '55,345,13,19,1,3,0,1\\n',\n",
       " '55,313,10,20,1,1,0,0\\n',\n",
       " '53,320,14,13,1,3,0,1\\n',\n",
       " '232,236,7,14,1,1,0,1\\n',\n",
       " '237,236,6,14,1,1,0,1\\n',\n",
       " '193,235,7,13,1,2,0,1\\n',\n",
       " '192,241,11,11,1,10,0,1\\n',\n",
       " '178,234,6,10,1,2,0,2\\n',\n",
       " '177,237,8,12,1,10,0,1\\n',\n",
       " '117,202,7,11,1,2,0,1\\n',\n",
       " '118,206,7,12,1,10,0,1\\n',\n",
       " '59,211,9,13,1,1,0,0\\n',\n",
       " '86,194,8,17,1,1,0,2\\n',\n",
       " '219,118,7,10,1,1,0,0\\n',\n",
       " '193,99,6,10,1,10,0,2\\n',\n",
       " '137,115,6,9,1,2,0,1\\n',\n",
       " '109,107,4,7,1,2,0,1\\n',\n",
       " '109,110,5,8,1,10,0,1\\n',\n",
       " '87,102,9,13,1,8,0,1\\n',\n",
       " '112,100,4,9,1,2,0,1\\n',\n",
       " '59,280,11,18,1,1,0,0\\n',\n",
       " '89,81,10,10,1,5,1,0\\n',\n",
       " '136,119,6,8,1,10,0,1\\n',\n",
       " '841,487,21,26,1,10,0,0\\n',\n",
       " '912,139,34,73,0,0,0,0\\n',\n",
       " '223,345,8,11,1,2,0,2\\n',\n",
       " '329,161,7,11,1,2,0,0\\n']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# visdrone data format\n",
    "# <bbox_left>,<bbox_top>,<bbox_width>,<bbox_height>,<score>,<object_category>,<truncation>,<occlusion>\n",
    "# 即：左上角坐标，宽，高，置信度，类别，截断，遮挡\n",
    "\n",
    "lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c90cd4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 目标YOLO格式：类别id、x_center y_center width height"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "55f081eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import glob\n",
    "import cv2\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7e26b72d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 原数据集的类别标签\n",
    "label_map = {\n",
    "    0: \"pedestrian\",\n",
    "    1: \"people\",\n",
    "    2: \"bicycle\",\n",
    "    3: \"car\",\n",
    "    4: \"van\",\n",
    "    5: \"truck\",\n",
    "    6: \"tricycle\",\n",
    "    7: \"awning-tricycle\",\n",
    "    8: \"bus\",\n",
    "    9: \"motor\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d27cd6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "根据原始标注文件、图片文件，转换为yolo格式的标注文件\n",
    "@data_file: 原始标注文件\n",
    "@img_file: 图片文件\n",
    "\n",
    "@results: yolo格式的标注文件\n",
    "\"\"\"\n",
    "def transform_data(data_file, img_file):\n",
    "    with open(data_file, \"r\", encoding=\"utf-8\") as fd:\n",
    "        lines = fd.readlines()\n",
    "\n",
    "    image = cv2.imread(img_file)\n",
    "    height, width, _ = image.shape\n",
    "\n",
    "    results = []\n",
    "    for line in lines:\n",
    "        element = line.strip().split(\",\")\n",
    "        bbox_left = float(element[0])\n",
    "        bbox_top = float(element[1])\n",
    "        bbox_width = float(element[2])\n",
    "        bbox_height = float(element[3])\n",
    "        object_category = int(element[5]) - 1\n",
    "\n",
    "        if object_category not in label_map.keys():\n",
    "            continue\n",
    "\n",
    "        center_x = (bbox_left + bbox_width / 2) / width\n",
    "        center_y = (bbox_top + bbox_height / 2) / height\n",
    "        b_width = bbox_width / width\n",
    "        b_height = bbox_height / height\n",
    "        results.append(\n",
    "            {\n",
    "                \"label\": object_category,\n",
    "                \"x\": center_x,\n",
    "                \"y\": center_y,\n",
    "                \"width\": b_width,\n",
    "                \"height\": b_height,\n",
    "            }\n",
    "        )\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44c03e12",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "将原始数据集转换为yolo格式\n",
    "@in_path: 原始数据集路径\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def process_dir(in_path):\n",
    "    img_path = os.path.join(in_path, \"images\")\n",
    "    ana_path = os.path.join(in_path, \"annotations\")\n",
    "\n",
    "    svd_path = os.path.join(in_path, \"labels\")\n",
    "    if not os.path.exists(svd_path):\n",
    "        os.makedirs(svd_path)\n",
    "\n",
    "    txt_files = sorted(glob.glob(os.path.join(ana_path, \"*.txt\")))\n",
    "    for txt_file in tqdm.tqdm(txt_files, desc=\"Processing\"):\n",
    "        base_name = os.path.splitext(os.path.basename(txt_file))[0]\n",
    "        img_file = os.path.join(img_path, base_name + \".jpg\")\n",
    "        if not os.path.exists(img_file):\n",
    "            # print(f\"{img_file} is not found\")\n",
    "            continue\n",
    "\n",
    "        result = transform_data(txt_file, img_file)\n",
    "\n",
    "        saved_file = os.path.join(svd_path, os.path.basename(txt_file))\n",
    "        with open(saved_file, \"w\", encoding=\"utf-8\") as fw:\n",
    "            for line in result:\n",
    "                # print(line)\n",
    "                write_line = (\n",
    "                    str(line[\"label\"])\n",
    "                    + \" \"\n",
    "                    + str(line[\"x\"])\n",
    "                    + \" \"\n",
    "                    + str(line[\"y\"])\n",
    "                    + \" \"\n",
    "                    + str(line[\"width\"])\n",
    "                    + \" \"\n",
    "                    + str(line[\"height\"])\n",
    "                    + \"\\n\"\n",
    "                )\n",
    "                fw.write(write_line)\n",
    "\n",
    "    print(\"Done...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a86dc088",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 处理train_dir\n",
    "process_dir(train_dir)\n",
    "\n",
    "# 处理val_dir\n",
    "process_dir(val_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4ace768",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f6572c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "根据yolo格式的标注文件，在图片上绘制\n",
    "@param img: 图片\n",
    "@param yoloLabelFile: yolo格式的标注文件路径\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def yoloDraw(img, yoloLabelFile):\n",
    "    img_copy = cv2.imread(img)\n",
    "    if img_copy is None:\n",
    "        print (\"Error: image is None\")\n",
    "        return\n",
    "    # 生成10类标签对应的颜色\n",
    "    color_dict = {\n",
    "        0: (255, 0, 0),\n",
    "        1: (0, 255, 0),\n",
    "        2: (0, 0, 255),\n",
    "        3: (255, 255, 0),\n",
    "        4: (0, 255, 255),\n",
    "        5: (255, 0, 255),\n",
    "        6: (255, 255, 255),\n",
    "        7: (0, 0, 0),\n",
    "        8: (128, 128, 128),\n",
    "        9: (128, 0, 0),\n",
    "    }\n",
    "    with open(yoloLabelFile, \"r\") as f:\n",
    "        lines = f.readlines()\n",
    "        lines = [line.strip() for line in lines]\n",
    "        boxes = [line.split(\" \") for line in lines]\n",
    "        for box in boxes:\n",
    "            class_label = int(box[0])\n",
    "            x_center, y_center, width, height = [float(i) for i in box[1:]]\n",
    "            x1 = int((x_center - width / 2) * img_copy.shape[1])\n",
    "            y1 = int((y_center - height / 2) * img_copy.shape[0])\n",
    "            x2 = int((x_center + width / 2) * img_copy.shape[1])\n",
    "            y2 = int((y_center + height / 2) * img_copy.shape[0])\n",
    "            cv2.rectangle(img_copy, (x1, y1), (x2, y2), color_dict[class_label], 2)\n",
    "            cv2.putText(\n",
    "                img_copy,\n",
    "                str(class_label),\n",
    "                (x1, y1),\n",
    "                cv2.FONT_HERSHEY_SIMPLEX,\n",
    "                1,\n",
    "                color_dict[class_label],\n",
    "                2,\n",
    "            )\n",
    "    plt.imshow(img_copy[:, :, ::-1])\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67233ab7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 训练集中选择图片及对应标注\n",
    "test_file = './VisDrone2019-DET-train/images/0000002_00005_d_0000014.jpg'\n",
    "label_file = './VisDrone2019-DET-train/labels/0000002_00005_d_0000014.txt'\n",
    "yoloDraw(test_file, label_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be538203",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 测试集中选择图片及对应标注\n",
    "test_file = './VisDrone2019-DET-val/images/0000001_04527_d_0000008.jpg'\n",
    "label_file = './VisDrone2019-DET-val/labels/0000001_04527_d_0000008.txt'\n",
    "yoloDraw(test_file, label_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e1d762",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51afbe0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
